import torch
from torch.utils import data
from dataset.dataset import CrackDataSet
from config.Load_DefaultConfig import DefaultConfig
from tqdm import tqdm

if __name__ == '__main__':

    opt = DefaultConfig()
    dataset = CrackDataSet(opt.data_set, dataset_choose="test")
    data_loader = data.DataLoader(
        dataset,
        batch_size=opt.batch_size,
        shuffle=opt.shuffle,
        num_workers=opt.num_workers
    )
    device = torch.device(opt.device if torch.cuda.is_available() else "cpu")  # 模型
    confusion_matrix = {"TP": 0, "FP": 0, "TN": 0, "FN": 0, "ALL": 0}
    best_epoch = {"name": None, "acc": 0, "pre": 0, "rec": 0, "FMean": 0}
    for epoch in range(0, 200, 1):
        model = torch.load(f"./checkpoint/epoch{epoch}.pth").eval()
        tqdm_bar = tqdm(enumerate(data_loader), desc=f"eval：epoch{epoch}", total=len(data_loader))
        for ii, (img, label) in tqdm_bar:
            img = img.to(device)
            model = model.to(device)
            label = label.cpu()

            out = model(img)
            out_label = out.argmax(dim=1).cpu()

            for la, out_la in zip(label, out_label):
                if la == 0:
                    if out_la == 0:
                        confusion_matrix["TP"] += 1
                    elif out_la == 1:
                        confusion_matrix["FN"] += 1
                elif la == 1:
                    if out_la == 0:
                        confusion_matrix["FP"] += 1
                    elif out_la == 1:
                        confusion_matrix["TN"] += 1
                confusion_matrix["ALL"] += 1
            try:
                acc = (confusion_matrix["TP"] + confusion_matrix["TN"]) / (confusion_matrix["ALL"])
                pre = (confusion_matrix["TP"]) / (confusion_matrix["TP"] + confusion_matrix["FP"])
                rec = (confusion_matrix["TP"]) / (confusion_matrix["TP"] + confusion_matrix["FN"])
                FMean = 2 * (pre * rec) / (pre + rec)
                tqdm_bar.set_postfix_str({
                    "acc": f"{acc:.6f}",
                    "pre": f"{pre:.6f}",
                    "rec": f"{rec:.6f}",
                    "FMean": f"{FMean:.6f}",
                })
            except:
                pass

        acc = (confusion_matrix["TP"] + confusion_matrix["TN"])/(confusion_matrix["ALL"])
        pre = (confusion_matrix["TP"])/(confusion_matrix["TP"] + confusion_matrix["FP"])
        rec = (confusion_matrix["TP"])/(confusion_matrix["TP"] + confusion_matrix["FN"])
        FMean = 2 * (pre * rec) / (pre + rec)
        if FMean > best_epoch["FMean"]:
            best_epoch["name"] = f"epoch{epoch}"
            best_epoch["acc"] = acc
            best_epoch["pre"] = pre
            best_epoch["rec"] = rec
            best_epoch["FMean"] = FMean
    print(best_epoch)


